import asyncio
import json
import logging
from typing import Dict, Optional

import websockets
from pydantic import BaseModel

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


# 定义MCP消息模型
class MCPMessage(BaseModel):
    type: str
    content: Dict
    id: Optional[str] = None


class MCPTool(BaseModel):
    name: str
    description: str
    parameters: Dict


# Ollama API配置
OLLAMA_API_URL = "http://localhost:11434/api"
MODEL_NAME = "qwen:8b"  # 修正模型名称

# 示例工具定义
TOOLS = [
    MCPTool(
        name="calculator",
        description="A simple calculator that can perform basic arithmetic operations",
        parameters={
            "type": "object",
            "properties": {
                "operation": {
                    "type": "string",
                    "enum": ["add", "subtract", "multiply", "divide"]
                },
                "a": {"type": "number"},
                "b": {"type": "number"}
            },
            "required": ["operation", "a", "b"]
        }
    ),
    MCPTool(
        name="weather",
        description="Get current weather information for a location",
        parameters={
            "type": "object",
            "properties": {
                "location": {"type": "string"},
                "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]}
            },
            "required": ["location"]
        }
    )
]


class MCPServer:
    def __init__(self):
        self.clients = set()
        self.ollama_connected = False

    async def check_ollama_connection(self):
        """检查Ollama服务是否可用"""
        try:
            async with websockets.connect(f"{OLLAMA_API_URL}/chat") as websocket:
                await websocket.send(json.dumps({
                    "model": MODEL_NAME,
                    "messages": [{"role": "user", "content": "test"}],
                    "stream": False
                }))
                response = await websocket.recv()
                self.ollama_connected = True
                logger.info("Successfully connected to Ollama service")
                return True
        except Exception as e:
            logger.error(f"Failed to connect to Ollama service: {e}")
            self.ollama_connected = False
            return False

    async def handle_client(self, websocket):
        """处理客户端连接"""
        client_id = id(websocket)
        self.clients.add(websocket)
        logger.info(f"New client connected: {client_id}")

        try:
            # 检查Ollama连接
            if not self.ollama_connected:
                if not await self.check_ollama_connection():
                    await websocket.send(json.dumps({
                        "type": "error",
                        "content": {"message": "Ollama service is not available"}
                    }))
                    return

            # 发送工具定义
            tools_message = {
                "type": "tools",
                "content": {"tools": [tool.dict() for tool in TOOLS]}
            }
            logger.debug(f"Sending tools definition: {json.dumps(tools_message)}")
            await websocket.send(json.dumps(tools_message))

            async for message in websocket:
                try:
                    message_data = json.loads(message)
                    logger.debug(f"Received message from client {client_id}: {message}")

                    if message_data["type"] == "chat":
                        # 调用Ollama API
                        try:
                            response = await self.call_ollama(message_data["content"]["message"])
                            await websocket.send(json.dumps({
                                "type": "response",
                                "content": {"message": response}
                            }))
                            logger.debug(f"Sent response to client {client_id}")
                        except Exception as e:
                            logger.error(f"Error calling Ollama: {e}")
                            await websocket.send(json.dumps({
                                "type": "error",
                                "content": {"message": f"Error calling Ollama: {str(e)}"}
                            }))
                    elif message_data["type"] == "tool_call":
                        # 处理工具调用
                        try:
                            tool_response = await self.handle_tool_call(message_data["content"])
                            await websocket.send(json.dumps({
                                "type": "tool_response",
                                "content": tool_response
                            }))
                            logger.debug(f"Sent tool response to client {client_id}")
                        except Exception as e:
                            logger.error(f"Error handling tool call: {e}")
                            await websocket.send(json.dumps({
                                "type": "error",
                                "content": {"message": f"Error handling tool call: {str(e)}"}
                            }))

                except json.JSONDecodeError as e:
                    logger.error(f"Invalid JSON message from client {client_id}: {e}")
                    await websocket.send(json.dumps({
                        "type": "error",
                        "content": {"message": "Invalid JSON message"}
                    }))
                except Exception as e:
                    logger.error(f"Error handling message from client {client_id}: {e}")
                    await websocket.send(json.dumps({
                        "type": "error",
                        "content": {"message": f"Server error: {str(e)}"}
                    }))

        except websockets.exceptions.ConnectionClosed as e:
            logger.info(f"Client {client_id} disconnected: {e}")
        except Exception as e:
            logger.error(f"Unexpected error with client {client_id}: {e}")
        finally:
            self.clients.remove(websocket)
            logger.info(f"Client {client_id} removed from active clients")

    async def call_ollama(self, message: str) -> str:
        """调用Ollama API进行对话"""
        if not self.ollama_connected:
            raise ConnectionError("Ollama service is not available")

        try:
            async with websockets.connect(f"{OLLAMA_API_URL}/chat") as websocket:
                request = {
                    "model": MODEL_NAME,
                    "messages": [{"role": "user", "content": message}],
                    "stream": False
                }
                logger.debug(f"Sending request to Ollama: {json.dumps(request)}")
                await websocket.send(json.dumps(request))

                response = await websocket.recv()
                response_data = json.loads(response)
                logger.debug(f"Received response from Ollama: {response}")

                if "message" not in response_data:
                    raise ValueError("Invalid response from Ollama")

                return response_data["message"]["content"]
        except Exception as e:
            logger.error(f"Error in call_ollama: {e}")
            self.ollama_connected = False
            raise

    async def handle_tool_call(self, tool_data: Dict) -> Dict:
        """处理工具调用"""
        tool_name = tool_data["name"]
        params = tool_data["parameters"]
        logger.debug(f"Handling tool call: {tool_name} with params: {params}")

        try:
            if tool_name == "calculator":
                return await self.handle_calculator(params)
            elif tool_name == "weather":
                return await self.handle_weather(params)
            else:
                logger.warning(f"Unknown tool called: {tool_name}")
                return {"error": "Unknown tool"}
        except Exception as e:
            logger.error(f"Error handling tool call {tool_name}: {e}")
            raise

    async def handle_calculator(self, params: Dict) -> Dict:
        """处理计算器工具调用"""
        operation = params["operation"]
        a = params["a"]
        b = params["b"]
        logger.debug(f"Calculator operation: {operation} with {a} and {b}")

        try:
            result = None
            if operation == "add":
                result = a + b
            elif operation == "subtract":
                result = a - b
            elif operation == "multiply":
                result = a * b
            elif operation == "divide":
                result = a / b if b != 0 else "Error: Division by zero"
            else:
                raise ValueError(f"Unknown operation: {operation}")

            return {
                "result": result,
                "operation": operation,
                "parameters": params
            }
        except Exception as e:
            logger.error(f"Error in calculator: {e}")
            raise

    async def handle_weather(self, params: Dict) -> Dict:
        """处理天气工具调用（模拟）"""
        location = params["location"]
        unit = params.get("unit", "celsius")
        logger.debug(f"Weather request for {location} in {unit}")

        try:
            # 模拟天气数据
            return {
                "location": location,
                "temperature": 25 if unit == "celsius" else 77,
                "unit": unit,
                "condition": "sunny",
                "humidity": 65
            }
        except Exception as e:
            logger.error(f"Error in weather: {e}")
            raise

    async def start(self, host: str = "localhost", port: int = 8000):
        """启动MCP服务器"""
        # 检查Ollama连接
        if not await self.check_ollama_connection():
            logger.error("Failed to connect to Ollama service. Server will not start.")
            return

        async with websockets.serve(self.handle_client, host, port):
            logger.info(f"MCP Server running at ws://{host}:{port}")
            try:
                await asyncio.Future()  # 运行直到被取消
            except asyncio.CancelledError:
                logger.info("Server shutdown initiated")
            except Exception as e:
                logger.error(f"Server error: {e}")
            finally:
                # 清理所有客户端连接
                for client in self.clients.copy():
                    try:
                        await client.close()
                    except:
                        pass
                self.clients.clear()
                logger.info("Server stopped")


if __name__ == "__main__":
    server = MCPServer()
    try:
        asyncio.run(server.start())
    except KeyboardInterrupt:
        logger.info("\nServer stopped by user")
    except Exception as e:
        logger.error(f"Server error: {e}")
